import torch
import torch.nn as nn
import torch.nn.functional as F
import strip_pooling as sp

# 定义 StripPooling 模块（代码同上，不再重复）

# 定义一个简单的 UNet 模型
class UNetWithStripPooling(nn.Module):
    def __init__(self, num_classes):
        super(UNetWithStripPooling, self).__init__()
        self.enc1 = self.conv_block(3, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)

        self.pool = nn.MaxPool2d(2, 2)

        self.center = self.conv_block(512, 1024)

        # 在中心位置添加 StripPooling 模块
        self.strip_pool = sp.StripPooling(
            1024, (20, 12), nn.BatchNorm2d, {'mode': 'bilinear', 'align_corners': True})

        self.up4 = self.up_conv(1024, 512)
        self.dec4 = self.conv_block(1024, 512)
        self.up3 = self.up_conv(512, 256)
        self.dec3 = self.conv_block(512, 256)
        self.up2 = self.up_conv(256, 128)
        self.dec2 = self.conv_block(256, 128)
        self.up1 = self.up_conv(128, 64)
        self.dec1 = self.conv_block(128, 64)

        self.final = nn.Conv2d(64, num_classes, kernel_size=1)

    def conv_block(self, in_channels, out_channels):
        block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        return block

    def up_conv(self, in_channels, out_channels):
        up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
        return up

    def forward(self, x):
        # 编码器
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))

        # 中心部分
        center = self.center(self.pool(e4))
        center = self.strip_pool(center)  # 应用 StripPooling 模块

        # 解码器
        d4 = self.up4(center)
        d4 = torch.cat((e4, d4), dim=1)
        d4 = self.dec4(d4)

        d3 = self.up3(d4)
        d3 = torch.cat((e3, d3), dim=1)
        d3 = self.dec3(d3)

        d2 = self.up2(d3)
        d2 = torch.cat((e2, d2), dim=1)
        d2 = self.dec2(d2)

        d1 = self.up1(d2)
        d1 = torch.cat((e1, d1), dim=1)
        d1 = self.dec1(d1)

        out = self.final(d1)
        return out


# 测试模型
if __name__ == '__main__':
    DEVICE = 'mps' if torch.cuda.is_available() else 'cpu'
    model = UNetWithStripPooling(num_classes=21).to(DEVICE)
    input = torch.randn(1, 3, 256, 256).to(DEVICE)
    output = model(input)
    print(input.size(), output.size())
    print(f"Output shape: {output.shape}")